"use client";

import { Button } from "@/components/ui/button";
import { memo, useCallback, useEffect, useRef, useState } from "react";
import Image from "next/image";
import { FixedSizeGrid as Grid, GridChildComponentProps } from "react-window";
import {
  Card,
  CardContent,
  CardHeader,
  CardTitle,
  CardDescription,
} from "@/components/ui/card";
import { config } from "@/config";
import { Skeleton } from "@/components/ui/skeleton";
import { AlertTriangle, ImageOff, RefreshCw } from "lucide-react";
import { ImageItem, ModalType } from "@/lib/types";
import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group";

const IMAGE_SIZE = 200;
const IMAGE_MARGIN = 16;
const MAX_CONCURRENT_LOADS = 6;

const ALL_MODALITIES: ModalType[] = ["SAR", "IR", "RGB"];

// 工具函数
const fetchData = async (url: string) => {
  const res = await fetch(url);
  if (!res.ok) throw new Error(`HTTP ${res.status}`);
  return res.json();
};

// 图片单元格组件
const Cell = memo(
  ({ columnIndex, rowIndex, style, data }: GridChildComponentProps) => {
    const { images, columnCount, fetchImageBlob, onRetry } = data;
    const index = rowIndex * columnCount + columnIndex;
    const imageItem = images[index] as ImageItem | undefined;

    useEffect(() => {
      if (
        imageItem &&
        !imageItem.blobUrl &&
        !imageItem.error &&
        !imageItem.loading
      ) {
        fetchImageBlob(imageItem.image_path);
      }
    }, [imageItem, fetchImageBlob]);

    if (!imageItem) return <div style={style}></div>;

    return (
      <div style={style} className="flex items-center justify-center p-1">
        <Card className="w-full h-full overflow-hidden flex flex-col hover:shadow-md transition-shadow">
          <CardContent className="p-2 flex-grow flex items-center justify-center relative">
            {imageItem.error ? (
              <div className="flex flex-col items-center justify-center text-destructive">
                <ImageOff size={48} />
                <p className="text-xs mt-1">加载失败</p>
                <Button
                  size="sm"
                  variant="outline"
                  className="mt-2"
                  onClick={() => onRetry(imageItem.image_path)}
                >
                  <RefreshCw size={12} className="mr-1" />
                  重试
                </Button>
              </div>
            ) : imageItem.blobUrl ? (
              <Image
                src={imageItem.blobUrl}
                alt={imageItem.image_name}
                width={IMAGE_SIZE - 16}
                height={IMAGE_SIZE - 16}
                className="object-contain"
                unoptimized
              />
            ) : (
              <Skeleton className="w-full h-full" />
            )}
          </CardContent>
          <div className="p-2 border-t text-xs">
            <p className="truncate font-medium" title={imageItem.image_name}>
              {imageItem.image_name}
            </p>
            <p className="text-muted-foreground">{imageItem.modality}</p>
          </div>
        </Card>
      </div>
    );
  }
);
Cell.displayName = "Cell";

export default function DatasetPage() {
  // 状态管理
  const [datasetTypes, setDatasetTypes] = useState<string[]>([]);
  const [selectedDatasetType, setSelectedDatasetType] = useState<string | null>(
    null
  );
  const [allImages, setAllImages] = useState<ImageItem[]>([]);
  const [selectedModality, setSelectedModality] = useState<ModalType>("RGB");
  const [loadingDatasets, setLoadingDatasets] = useState(true);
  const [loadingImages, setLoadingImages] = useState(false);
  const [errorDatasets, setErrorDatasets] = useState<string | null>(null);
  const [errorImages, setErrorImages] = useState<string | null>(null);
  const [loadingCount, setLoadingCount] = useState(0);
  const [gridColumnCount, setGridColumnCount] = useState(6);
  const [gridDimensions, setGridDimensions] = useState({
    width: 1400,
    height: 700,
  });

  const gridContainerRef = useRef<HTMLDivElement>(null);

  // 过滤图片
  const filteredImages = allImages.filter(
    (img) => img.modality === selectedModality
  );

  // 清理blob URLs
  useEffect(() => {
    return () => {
      allImages.forEach((img) => {
        if (img.blobUrl) URL.revokeObjectURL(img.blobUrl);
      });
    };
  }, [allImages]);

  // 网格尺寸计算
  useEffect(() => {
    const updateDimensions = () => {
      if (gridContainerRef.current) {
        const rect = gridContainerRef.current.getBoundingClientRect();
        const width = rect.width || 800;
        const height = rect.height || 600;
        setGridDimensions({ width, height });
        setGridColumnCount(
          Math.max(1, Math.floor(width / (IMAGE_SIZE + IMAGE_MARGIN)))
        );
      }
    };

    updateDimensions();
    const resizeObserver = new ResizeObserver(updateDimensions);
    if (gridContainerRef.current) {
      resizeObserver.observe(gridContainerRef.current);
    }
    return () => resizeObserver.disconnect();
  }, []);

  // 获取数据集类型
  useEffect(() => {
    const loadDatasetTypes = async () => {
      try {
        const data = await fetchData(
          `${config.AI_SERVICE_URL}${config.AI_SERVICE_URL_PREFIX}/datasets`
        );
        setDatasetTypes(data);
        if (data.length > 0) setSelectedDatasetType(data[0]);
      } catch {
        setErrorDatasets("无法加载数据集类型");
      } finally {
        setLoadingDatasets(false);
      }
    };
    loadDatasetTypes();
  }, []);

  // 获取图片列表
  useEffect(() => {
    if (!selectedDatasetType) {
      setAllImages([]);
      return;
    }

    const loadImages = async () => {
      setLoadingImages(true);
      setErrorImages(null);

      // 清理之前的图片
      allImages.forEach((img) => {
        if (img.blobUrl) URL.revokeObjectURL(img.blobUrl);
      });
      setAllImages([]);

      try {
        const data = await fetchData(
          `${config.AI_SERVICE_URL}${config.AI_SERVICE_URL_PREFIX}/images/${selectedDatasetType}`
        );
        setAllImages(
          data.map((img: ImageData, index: number) => ({
            ...img,
            id: `${selectedDatasetType}-${index}`,
            loading: false,
          }))
        );
      } catch {
        setErrorImages(`无法加载 ${selectedDatasetType} 的图片`);
      } finally {
        setLoadingImages(false);
      }
    };

    loadImages();
  }, [selectedDatasetType]);

  // 获取单张图片
  const fetchImageBlob = useCallback(
    async (imagePath: string) => {
      if (loadingCount >= MAX_CONCURRENT_LOADS) {
        setTimeout(() => fetchImageBlob(imagePath), 500);
        return;
      }

      setLoadingCount((count) => count + 1);
      setAllImages((prev) =>
        prev.map((img) =>
          img.image_path === imagePath ? { ...img, loading: true } : img
        )
      );

      try {
        const res = await fetch(
          `${config.AI_SERVICE_URL}${
            config.AI_SERVICE_URL_PREFIX
          }/image?image_path=${encodeURIComponent(imagePath)}`
        );

        if (!res.ok) throw new Error(`HTTP ${res.status}`);

        const blob = await res.blob();
        const blobUrl = URL.createObjectURL(blob);

        setAllImages((prev) =>
          prev.map((img) =>
            img.image_path === imagePath
              ? { ...img, blobUrl, loading: false, error: false }
              : img
          )
        );
      } catch {
        setAllImages((prev) =>
          prev.map((img) =>
            img.image_path === imagePath
              ? { ...img, error: true, loading: false }
              : img
          )
        );
      } finally {
        setLoadingCount((count) => count - 1);
      }
    },
    [loadingCount]
  );

  // 重试单张图片
  const handleRetryImage = useCallback(
    (imagePath: string) => {
      setAllImages((prev) =>
        prev.map((img) =>
          img.image_path === imagePath
            ? { ...img, error: false, loading: false }
            : img
        )
      );
      fetchImageBlob(imagePath);
    },
    [fetchImageBlob]
  );

  // 事件处理
  const handleDatasetTypeClick = (type: string) => {
    setSelectedDatasetType(type);
    setSelectedModality("RGB");
  };

  const rowCount = Math.ceil(filteredImages.length / gridColumnCount);

  // 渲染逻辑
  if (loadingDatasets) {
    return (
      <div className="flex-1 h-full m-4 flex items-center justify-center">
        <div className="text-center">
          <div className="animate-spin w-8 h-8 border-2 border-primary border-t-transparent rounded-full mx-auto mb-2" />
          <p>正在加载数据集...</p>
        </div>
      </div>
    );
  }

  if (errorDatasets) {
    return (
      <div className="flex-1 h-full m-4 flex flex-col items-center justify-center text-destructive space-y-4">
        <AlertTriangle size={48} />
        <p>{errorDatasets}</p>
        <Button onClick={() => window.location.reload()}>重新加载</Button>
      </div>
    );
  }

  return (
    <div className="flex-1 h-full py-2 flex flex-col gap-2">
      {/* 数据集选择 */}
      <Card>
        <CardHeader>
          <CardTitle>选择数据集类型</CardTitle>
          <CardDescription>点击下方按钮加载对应数据集的图片</CardDescription>
        </CardHeader>
        <CardContent className="flex flex-wrap gap-2">
          {datasetTypes.map((type) => (
            <Button
              key={type}
              variant={selectedDatasetType === type ? "default" : "secondary"}
              onClick={() => handleDatasetTypeClick(type)}
              disabled={loadingImages}
            >
              {type}
            </Button>
          ))}
        </CardContent>
      </Card>

      {/* 模态筛选 */}
      {selectedDatasetType && (
        <Card>
          <CardHeader>
            <CardTitle>筛选图片模态</CardTitle>
            <CardDescription>
              显示 {filteredImages.length} / {allImages.length} 张图片
            </CardDescription>
          </CardHeader>
          <CardContent className="flex flex-wrap gap-4">
            <RadioGroup
              value={selectedModality}
              onValueChange={(value) => setSelectedModality(value as ModalType)}
              className="flex flex-wrap gap-4"
            >
              {ALL_MODALITIES.map((modality) => (
                <div key={modality} className="flex items-center space-x-2">
                  <RadioGroupItem
                    id={`modality-${modality}`}
                    value={modality}
                  />
                  <label
                    htmlFor={`modality-${modality}`}
                    className="text-sm font-medium"
                  >
                    {modality}
                  </label>
                </div>
              ))}
            </RadioGroup>
          </CardContent>
        </Card>
      )}

      {/* 加载状态 */}
      {loadingImages && (
        <div className="flex-1 flex items-center justify-center">
          <div className="text-center">
            <div className="animate-spin w-8 h-8 border-2 border-primary border-t-transparent rounded-full mx-auto mb-2" />
            <p>正在加载图片列表...</p>
          </div>
        </div>
      )}

      {/* 错误状态 */}
      {errorImages && (
        <div className="flex-1 flex flex-col items-center justify-center text-destructive space-y-4">
          <AlertTriangle size={48} />
          <p>{errorImages}</p>
          <Button onClick={() => handleDatasetTypeClick(selectedDatasetType!)}>
            重试
          </Button>
        </div>
      )}

      {/* 无图片提示 */}
      {!loadingImages &&
        !errorImages &&
        selectedDatasetType &&
        filteredImages.length === 0 && (
          <div className="flex-1 flex items-center justify-center">
            <p className="text-muted-foreground">当前条件下没有图片可显示</p>
          </div>
        )}

      {/* 图片网格 */}
      {!loadingImages &&
        !errorImages &&
        selectedDatasetType &&
        filteredImages.length > 0 && (
          <div ref={gridContainerRef} className="flex-1 w-full overflow-auto">
            <Grid
              columnCount={gridColumnCount}
              columnWidth={IMAGE_SIZE + IMAGE_MARGIN}
              height={gridDimensions.height}
              rowCount={rowCount}
              rowHeight={IMAGE_SIZE + IMAGE_MARGIN + 60}
              width={gridDimensions.width}
              itemData={{
                images: filteredImages,
                columnCount: gridColumnCount,
                fetchImageBlob,
                onRetry: handleRetryImage,
              }}
              overscanRowCount={2}
            >
              {Cell}
            </Grid>
          </div>
        )}

      {/* 初始状态 */}
      {!selectedDatasetType && !loadingDatasets && !errorDatasets && (
        <div className="flex-1 flex items-center justify-center">
          <p className="text-muted-foreground">
            请先选择一个数据集类型以查看图片
          </p>
        </div>
      )}
    </div>
  );
}
